#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
JSON编码工具类
解决JSON数据传输过程中的编码错误问题
"""

import json
import logging
import base64
import struct
from typing import Dict, Any, Optional, Union, List
from datetime import datetime
import unicodedata

logger = logging.getLogger(__name__)

class JSONEncoder:
    """JSON编码器 - 处理编码错误和特殊字符"""
    
    def __init__(self, encoding='utf-8', ensure_ascii=False, default_encoding='utf-8'):
        self.encoding = encoding
        self.ensure_ascii = ensure_ascii
        self.default_encoding = default_encoding
    
    def clean_string(self, text: str) -> str:
        """清理字符串，移除或替换问题字符"""
        if not isinstance(text, str):
            return str(text)
        
        # 移除控制字符（除了换行符和制表符）
        cleaned = ''
        for char in text:
            if unicodedata.category(char)[0] == 'C':
                # 控制字符
                if char in '\n\t\r':
                    cleaned += char
                else:
                    cleaned += ' '  # 替换为空格
            else:
                cleaned += char
        
        # 移除零宽字符
        cleaned = cleaned.replace('\u200b', '')  # 零宽空格
        cleaned = cleaned.replace('\u200c', '')  # 零宽非连接符
        cleaned = cleaned.replace('\u200d', '')  # 零宽连接符
        cleaned = cleaned.replace('\ufeff', '')  # BOM
        
        return cleaned.strip()
    
    def clean_dict(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """递归清理字典中的所有字符串"""
        cleaned_data = {}
        
        for key, value in data.items():
            # 清理键名
            clean_key = self.clean_string(str(key))
            
            # 清理值
            if isinstance(value, str):
                clean_value = self.clean_string(value)
            elif isinstance(value, dict):
                clean_value = self.clean_dict(value)
            elif isinstance(value, list):
                clean_value = self.clean_list(value)
            elif isinstance(value, datetime):
                clean_value = value.isoformat()
            else:
                clean_value = value
            
            cleaned_data[clean_key] = clean_value
        
        return cleaned_data
    
    def clean_list(self, data: List[Any]) -> List[Any]:
        """递归清理列表中的所有字符串"""
        cleaned_list = []
        
        for item in data:
            if isinstance(item, str):
                clean_item = self.clean_string(item)
            elif isinstance(item, dict):
                clean_item = self.clean_dict(item)
            elif isinstance(item, list):
                clean_item = self.clean_list(item)
            elif isinstance(item, datetime):
                clean_item = item.isoformat()
            else:
                clean_item = item
            
            cleaned_list.append(clean_item)
        
        return cleaned_list
    
    def safe_json_dumps(self, data: Any, **kwargs) -> str:
        """安全的JSON序列化，处理编码错误"""
        try:
            # 清理数据
            if isinstance(data, dict):
                cleaned_data = self.clean_dict(data)
            elif isinstance(data, list):
                cleaned_data = self.clean_list(data)
            else:
                cleaned_data = data
            
            # 设置默认参数
            default_kwargs = {
                'ensure_ascii': self.ensure_ascii,
                'indent': None,
                'separators': (',', ':'),
                'default': self.json_default
            }
            default_kwargs.update(kwargs)
            
            # 尝试序列化
            json_str = json.dumps(cleaned_data, **default_kwargs)
            return json_str
            
        except UnicodeEncodeError as e:
            logger.warning(f"Unicode编码错误，尝试修复: {e}")
            # 强制使用ASCII编码
            try:
                return json.dumps(cleaned_data, ensure_ascii=True, **kwargs)
            except Exception as e2:
                logger.error(f"修复编码错误失败: {e2}")
                raise
        except Exception as e:
            logger.error(f"JSON序列化失败: {e}")
            raise
    
    def json_default(self, obj):
        """JSON默认序列化器"""
        if isinstance(obj, datetime):
            return obj.isoformat()
        elif hasattr(obj, 'isoformat'):
            return obj.isoformat()
        elif hasattr(obj, '__dict__'):
            return obj.__dict__
        else:
            return str(obj)
    
    def safe_json_loads(self, json_str: str, **kwargs) -> Any:
        """安全的JSON反序列化，处理编码错误"""
        try:
            return json.loads(json_str, **kwargs)
        except UnicodeDecodeError as e:
            logger.warning(f"Unicode解码错误，尝试修复: {e}")
            # 尝试不同的编码
            for encoding in ['utf-8', 'gbk', 'gb2312', 'latin1']:
                try:
                    if isinstance(json_str, bytes):
                        decoded_str = json_str.decode(encoding)
                    else:
                        decoded_str = json_str
                    return json.loads(decoded_str, **kwargs)
                except (UnicodeDecodeError, json.JSONDecodeError):
                    continue
            raise
        except json.JSONDecodeError as e:
            logger.error(f"JSON解析错误: {e}")
            raise

class JSONProtocol:
    """JSON协议处理器 - 处理网络传输协议"""
    
    def __init__(self, encoder: JSONEncoder = None):
        self.encoder = encoder or JSONEncoder()
        self.max_message_size = 10 * 1024 * 1024  # 10MB
    
    def encode_message(self, data: Dict[str, Any]) -> bytes:
        """编码JSON消息，添加长度前缀"""
        try:
            # 序列化JSON
            json_str = self.encoder.safe_json_dumps(data)
            json_bytes = json_str.encode(self.encoder.encoding)
            
            # 检查消息大小
            if len(json_bytes) > self.max_message_size:
                raise ValueError(f"消息太大: {len(json_bytes)} 字节，最大允许: {self.max_message_size} 字节")
            
            # 添加长度前缀 (4字节大端序)
            length_bytes = struct.pack('>I', len(json_bytes))
            message = length_bytes + json_bytes
            
            logger.debug(f"编码JSON消息: {len(json_bytes)} 字节")
            return message
            
        except Exception as e:
            logger.error(f"编码JSON消息失败: {e}")
            raise
    
    def decode_message(self, data: bytes) -> Optional[Dict[str, Any]]:
        """解码JSON消息，处理长度前缀"""
        try:
            if len(data) < 4:
                return None  # 数据不足，等待更多数据
            
            # 读取长度前缀
            message_length = struct.unpack('>I', data[:4])[0]
            
            # 检查消息大小
            if message_length > self.max_message_size:
                raise ValueError(f"消息太大: {message_length} 字节，最大允许: {self.max_message_size} 字节")
            
            # 检查是否有完整的消息
            if len(data) < 4 + message_length:
                return None  # 数据不足，等待更多数据
            
            # 提取JSON数据
            json_bytes = data[4:4 + message_length]
            
            # 解码JSON
            json_str = json_bytes.decode(self.encoder.encoding)
            result = self.encoder.safe_json_loads(json_str)
            
            logger.debug(f"解码JSON消息: {message_length} 字节")
            return result
            
        except Exception as e:
            logger.error(f"解码JSON消息失败: {e}")
            raise
    
    def extract_complete_messages(self, buffer: bytes) -> tuple:
        """从缓冲区中提取完整的消息"""
        messages = []
        remaining_buffer = buffer
        
        while len(remaining_buffer) >= 4:
            try:
                # 读取长度前缀
                message_length = struct.unpack('>I', remaining_buffer[:4])[0]
                
                # 检查消息大小
                if message_length > self.max_message_size:
                    logger.error(f"消息太大，跳过: {message_length} 字节")
                    remaining_buffer = remaining_buffer[4:]
                    continue
                
                # 检查是否有完整的消息
                if len(remaining_buffer) < 4 + message_length:
                    break  # 数据不足，等待更多数据
                
                # 提取消息
                message_data = remaining_buffer[4:4 + message_length]
                message = self.decode_message(remaining_buffer[:4 + message_length])
                
                if message is not None:
                    messages.append(message)
                
                # 更新缓冲区
                remaining_buffer = remaining_buffer[4 + message_length:]
                
            except Exception as e:
                logger.error(f"提取消息失败: {e}")
                # 跳过当前字节，继续处理
                remaining_buffer = remaining_buffer[1:]
        
        return messages, remaining_buffer

class JSONErrorHandler:
    """JSON错误处理器"""
    
    @staticmethod
    def handle_encoding_error(data: Any, error: Exception) -> str:
        """处理编码错误"""
        logger.warning(f"处理编码错误: {error}")
        
        try:
            # 尝试强制转换为字符串
            if isinstance(data, dict):
                # 递归处理字典
                return JSONErrorHandler._handle_dict_encoding_error(data)
            elif isinstance(data, list):
                # 递归处理列表
                return JSONErrorHandler._handle_list_encoding_error(data)
            else:
                # 直接转换为字符串
                return str(data)
        except Exception as e:
            logger.error(f"处理编码错误失败: {e}")
            return "{}"
    
    @staticmethod
    def _handle_dict_encoding_error(data: Dict[str, Any]) -> str:
        """处理字典编码错误"""
        try:
            cleaned_data = {}
            for key, value in data.items():
                try:
                    # 清理键名
                    clean_key = str(key).encode('ascii', 'ignore').decode('ascii')
                    
                    # 清理值
                    if isinstance(value, str):
                        clean_value = value.encode('ascii', 'ignore').decode('ascii')
                    elif isinstance(value, (dict, list)):
                        clean_value = JSONErrorHandler.handle_encoding_error(value, Exception("Nested error"))
                    else:
                        clean_value = str(value)
                    
                    cleaned_data[clean_key] = clean_value
                except Exception as e:
                    logger.warning(f"跳过问题字段 {key}: {e}")
                    continue
            
            return json.dumps(cleaned_data, ensure_ascii=True)
        except Exception as e:
            logger.error(f"处理字典编码错误失败: {e}")
            return "{}"
    
    @staticmethod
    def _handle_list_encoding_error(data: List[Any]) -> str:
        """处理列表编码错误"""
        try:
            cleaned_list = []
            for item in data:
                try:
                    if isinstance(item, str):
                        clean_item = item.encode('ascii', 'ignore').decode('ascii')
                    elif isinstance(item, (dict, list)):
                        clean_item = JSONErrorHandler.handle_encoding_error(item, Exception("Nested error"))
                    else:
                        clean_item = str(item)
                    
                    cleaned_list.append(clean_item)
                except Exception as e:
                    logger.warning(f"跳过问题列表项: {e}")
                    continue
            
            return json.dumps(cleaned_list, ensure_ascii=True)
        except Exception as e:
            logger.error(f"处理列表编码错误失败: {e}")
            return "[]"

# 全局实例
default_encoder = JSONEncoder()
default_protocol = JSONProtocol(default_encoder)

def safe_json_dumps(data: Any, **kwargs) -> str:
    """安全的JSON序列化函数"""
    return default_encoder.safe_json_dumps(data, **kwargs)

def safe_json_loads(json_str: str, **kwargs) -> Any:
    """安全的JSON反序列化函数"""
    return default_encoder.safe_json_loads(json_str, **kwargs)

def encode_json_message(data: Dict[str, Any]) -> bytes:
    """编码JSON消息"""
    return default_protocol.encode_message(data)

def decode_json_message(data: bytes) -> Optional[Dict[str, Any]]:
    """解码JSON消息"""
    return default_protocol.decode_message(data)
